#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <cluster.h>

#include <vector>
#include <iostream>
#include <iomanip>

using namespace sds;
using namespace std;

int main(int argc, const char* argv[]) {
  SFile reffile(argv[2]);
  SDataView<unsigned char> ref =
    reffile.getDataset(0).getDataView<unsigned char>();

  vector<SInteger> size = ref.getSize();

  SFile srcfile(argv[1]);
  SDataView<Point> src(size);

  for (int k = 0; k < CH; ++k) {
    SDataView<unsigned char> view =
      srcfile.getDataset(k).getDataView<unsigned char>();
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	((Point&)src[i][j])[k] = (double)view[i][j];
      }    
    }
  }

  vector<Cluster> clusters(256);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ref[i][j] != 0) {
	clusters[ref[i][j]] += src[i][j];
      }
    }
  }

  cout << setprecision(4);
  for (int k = 0; k < 256; ++k) {
    if (clusters[k].size() == 0) continue;
    cout << "#" << k << ' ' << '(' << clusters[k].size() << ')' << endl;
    for (int i = 0; i < CH; ++i) {
      cout << clusters[k].expected(i) << '\t';
      cout << sqrt(clusters[k].covariance(i, i)) << '\t' << '\t';
      for (int j = 0; j < CH; ++j) {
	cout << clusters[k].covariance(i, j) << '\t';
      }
      cout << endl;
    }
    cout << endl;
  }

  SDataView<unsigned char> trg(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      int minPos = 0; 
      double minVal = 0.0;
      for (int k = 0; k < 256; ++k) {
	if (clusters[k].size() == 0) continue;
	double dist = distance_rp(clusters[k].center(), src[i][j]);
	if (minPos == 0 || minVal > dist) {
	  minPos = k;
	  minVal = dist;
	}
      }
      trg[i][j] = minPos;
    }
  }
  
  SFile trgfile(argv[3], SFile::CREATE);
  trgfile.createDataset("clasified", UINT8, size).
    setDataView<unsigned char>(trg);
  return 0;
}
